'''
Author: 梦付千秋星垂野 465943794@qq.com
Date: 2022-06-27 09:34:45
LastEditors: 梦付千秋星垂野 465943794@qq.com
LastEditTime: 2022-07-04 10:26:47
FilePath: /base_machinelearning/Smartcity_CNN_Fish/algorithm/train.py
Description: trian文件
'''

import os
import sys
path=os.getcwd()
print(path)
sys.path.append(path)
print(sys.path)
import numpy as np
import platform
import torch.optim as optim
import torch
import torch.utils.data as Data
from Smartcity_CNN_Fish.algorithm.connect_db import readData
from Smartcity_CNN_Fish.algorithm.CNN_model import MLP
list_feature = ['company', 'consume', 'capital', 'house_price','employees','population']
num_inputs=len(list_feature)
feature_y="GDP"
distinct="data_pudong"

def inverse_transform_col(scaler, y, n_col):
    y = y.copy()
    y -= scaler.min_[n_col]
    #print("scalerY.min_:",scaler.min_)
    y /= scaler.scale_[n_col]
    return y
'''
author: ytc
description: train the cnn model
return {*}
'''   
def train():
    torch.manual_seed(2)
    trainX,trainY,scalerX,scalerY=readData(list_feature,feature_y,distinct)
    print(trainX.shape)
    batch_size = 30    # 将训练数据的特征和标签组合
    dataset = Data.TensorDataset(torch.tensor(trainX),torch.tensor(trainY))
    # 随机读取小批量
    data_iter = Data.DataLoader(dataset, batch_size, shuffle=True)
    num_epochs=1000
    net=MLP(num_inputs)
    optimizer=torch.optim.Adam(net.parameters(),lr=0.05)
    loss_func=torch.nn.MSELoss()
    #print(data_iter)
    net.train()
    for epoch in range(1, num_epochs + 1):
        for X, y in data_iter:
            output = net(X)
            l = loss_func(output.float(), y.view(-1, 1).float())
            optimizer.zero_grad() # 梯度清零，等价于net.zero_grad()
            l.backward()
            optimizer.step()
        print('epoch %d, loss: %f' % (epoch, l.item()))
    net.eval()
    with torch.no_grad():
        input_test_data = torch.tensor([scalerX.transform([[400049,994.89,57203,99574,198.27,319.72],[2264,7.04,6587,563,5.01,251.93]])]).float()
        #input_test_data = torch.tensor(scalerX.transform([[400049,994.89,57203,99574,198.27,319.72]])).float()
        prediction=net(input_test_data)
        result=inverse_transform_col(scalerY,prediction.detach().numpy(), 0)
        print(result)
    path=""
    if platform.system().lower() == 'windows':
        path=".\\Smartcity_CNN_Fish\\weight\\model_num_epochs_{}.pt".format(num_epochs)
    elif platform.system().lower() == 'linux':
        path="./Smartcity_CNN_Fish/weight/model_num_epochs_{}.pt".format(num_epochs)
    #torch.save(net.state_dict(), ".\\Smartcity_CNN_Fish\\weight\\model_num_epochs_{}.pt".format(num_epochs)) # 推荐的文件后缀名是pt或pth
    torch.save(net.state_dict(),path)
def predict(input,PATH):
    net=MLP(num_inputs)
    net.load_state_dict(torch.load(PATH))
    trainX,trainY,scalerX,scalerY=readData(list_feature,feature_y,distinct)
    input=torch.tensor(scalerX.transform(input)).float()
    net.eval()
    with torch.no_grad():
        predict=net(input)
        # print(scalerX.min_/scalerX.scale_)
        predict=inverse_transform_col(scalerY,predict.detach().numpy(), 0)
        # print(predict.shape)
        # print(predict)
        return predict
if __name__=="__main__":
    train()
